# 数据采集模块
import os
import pandas as pd
import akshare as ak
import datetime
import json
import sys
import logging
sys.path.append('..')
from config import DATA_PATH

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

class DataCollector:
    def __init__(self):
        # 确保数据存储目录存在
        if not os.path.exists(DATA_PATH):
            os.makedirs(DATA_PATH)
            logger.info(f"创建数据存储目录: {DATA_PATH}")
    
    def get_stock_basic(self, stock_code):
        """获取股票基本信息"""
        try:
            # 验证股票代码格式
            if not stock_code or not stock_code.isdigit() or len(stock_code) != 6:
                error_msg = f"无效的股票代码: {stock_code}，请输入6位数字代码"
                logger.error(error_msg)
                return {'error': error_msg}

            logger.info(f"开始获取股票 {stock_code} 的基本信息")
            logger.debug(f"调用AKShare API: stock_individual_info_em，参数：symbol={stock_code}")
            
            df = ak.stock_individual_info_em(symbol=stock_code)
            
            if df is None or df.empty:
                error_msg = f"未获取到股票 {stock_code} 的数据，可能是股票代码不存在或已退市"
                logger.error(error_msg)
                return {'error': error_msg}

            try:
                # 转换为与原接口相似的格式
                result = {
                    'ts_code': f'{stock_code}.{"6":"SH","0":"SZ","3":"SZ"}[stock_code[0]]',
                    'symbol': stock_code,
                    'name': str(df.loc[df['item'] == '股票简称', 'value'].iloc[0]),
                    'area': '中国',  # AKShare接口不提供地区信息，默认为中国
                    'industry': str(df.loc[df['item'] == '行业', 'value'].iloc[0]),
                    'list_date': str(df.loc[df['item'] == '上市时间', 'value'].iloc[0])
                }
                logger.info(f"股票 {stock_code} 基本信息获取成功")
                logger.debug(f"处理后的数据: {result}")
                return result
            except Exception as e:
                error_msg = f"数据处理失败: {str(e)}"
                logger.error(error_msg)
                return {'error': error_msg}
        except Exception as e:
            error_msg = f"获取股票基本信息失败: {type(e).__name__} - {str(e)}"
            logger.error(error_msg)
            return {'error': error_msg}
    
    def get_daily_data(self, stock_code, start_date=None, end_date=None):
        """获取股票日K线数据"""
        if not start_date:
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y-%m-%d')
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y-%m-%d')
        
        try:
            logger.info(f"开始获取股票 {stock_code} 的日K线数据")
            logger.debug(f"请求参数: start_date={start_date}, end_date={end_date}")
            
            df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", start_date=start_date, end_date=end_date, adjust="qfq")
            logger.info(f"获取到 {len(df)} 条数据记录")
            
            if not df.empty:
                # 转换列名以匹配原接口格式
                df = df.rename(columns={
                    '日期': 'trade_date',
                    '开盘': 'open',
                    '收盘': 'close',
                    '最高': 'high',
                    '最低': 'low',
                    '成交量': 'vol',
                    '成交额': 'amount'
                })
                
                # 添加ts_code列
                df['ts_code'] = f'{stock_code}.{"6":"SH","0":"SZ","3":"SZ"}[stock_code[0]]'
                
                # 确保日期格式统一
                df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y%m%d')
                df = df.sort_values('trade_date')
                logger.info("数据处理完成")
                return df
            
            logger.warning(f"未获取到股票 {stock_code} 的K线数据")
            return pd.DataFrame()
        except Exception as e:
            logger.error(f"获取日K线数据失败: {type(e).__name__} - {str(e)}")
            logger.error(f"请求参数: stock_code={stock_code}, start_date={start_date}, end_date={end_date}")
            return pd.DataFrame()
    
    def get_index_data(self, index_code='000001', start_date=None, end_date=None):
        """获取指数数据，默认为上证指数"""
        if not start_date:
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y-%m-%d')
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y-%m-%d')
        
        try:
            # 使用AKShare获取指数数据
            df = ak.stock_zh_index_hist(symbol=f"sh{index_code}" if index_code.startswith('0') else f"sz{index_code}",
                                       period="daily",
                                       start_date=start_date,
                                       end_date=end_date)
            
            # 转换列名以匹配原接口格式
            df = df.rename(columns={
                '日期': 'trade_date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'vol',
                '成交额': 'amount'
            })
            
            # 过滤日期范围
            df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y%m%d')
            df = df.sort_values('trade_date')
            return df
        except Exception as e:
            logger.error(f"获取指数数据失败: {e}")
            return pd.DataFrame()
    
    def get_industry_data(self, industry_code, start_date=None, end_date=None):
        """获取行业指数数据"""
        if not start_date:
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        try:
            # 使用AKShare获取行业指数数据
            df = ak.stock_board_industry_hist_em(symbol=industry_code, start_date=start_date, end_date=end_date)
            
            # 转换列名以匹配原接口格式
            df = df.rename(columns={
                '日期': 'trade_date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'vol',
                '成交额': 'amount'
            })
            
            # 确保日期格式统一
            df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y%m%d')
            df = df.sort_values('trade_date')
            return df
        except Exception as e:
            logger.error(f"获取行业数据失败: {e}")
            return pd.DataFrame()
    
    def get_money_flow(self, stock_code, start_date=None, end_date=None):
        """获取股票资金流向数据"""
        if not start_date:
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        try:
            # 根据股票代码判断市场类型
            market = "sh" if stock_code.startswith('6') else "sz" if stock_code.startswith(('0', '3')) else "bj"
            
            # 使用AKShare获取资金流向数据
            df = ak.stock_individual_fund_flow(stock=stock_code, market=market)
            
            if df is None or df.empty:
                logger.warning(f"获取股票 {stock_code} 的资金流向数据失败")
                return pd.DataFrame()
            
            # 转换日期格式
            df['trade_date'] = pd.to_datetime(df['日期']).dt.strftime('%Y%m%d')
            
            # 转换列名以匹配原接口格式
            df = df.rename(columns={
                '主力净流入-净额': 'net_mf_amount',
                '主力净流入-净占比': 'net_mf_amount_rate',
                '超大单净流入-净额': 'buy_lg_amount',
                '超大单净流入-净占比': 'buy_lg_rate',
                '大单净流入-净额': 'buy_md_amount',
                '大单净流入-净占比': 'buy_md_rate',
                '中单净流入-净额': 'buy_sm_amount',
                '中单净流入-净占比': 'buy_sm_rate',
                '小单净流入-净额': 'buy_xs_amount',
                '小单净流入-净占比': 'buy_xs_rate'
            })
            
            # 过滤日期范围
            if start_date and end_date:
                df = df[(df['trade_date'] >= start_date) & (df['trade_date'] <= end_date)]
            
            df = df.sort_values('trade_date')
            return df
        except Exception as e:
            logger.error(f"获取资金流向数据失败: {e}")
            return pd.DataFrame()
    
    def get_news(self, stock_code, limit=10):
        """获取股票相关新闻（使用Akshare）"""
        try:
            # 移除可能的后缀
            if '.' in stock_code:
                stock_code = stock_code.split('.')[0]
                
            # 使用东方财富网股票新闻接口
            news_df = ak.stock_news_em(symbol=stock_code)
            return news_df.head(limit) if not news_df.empty else pd.DataFrame()
        except Exception as e:
            print(f"获取新闻数据失败: {e}")
            return pd.DataFrame()
    
    def save_data(self, data, filename):
        """保存数据到本地"""
        file_path = os.path.join(DATA_PATH, filename)
        try:
            if isinstance(data, pd.DataFrame):
                data.to_csv(file_path, index=False)
            else:
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(data, f, ensure_ascii=False)
            return True
        except Exception as e:
            print(f"保存数据失败: {e}")
            return False
    
    def load_data(self, filename):
        """从本地加载数据"""
        file_path = os.path.join(DATA_PATH, filename)
        try:
            if filename.endswith('.csv'):
                return pd.read_csv(file_path)
            else:
                with open(file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
        except Exception as e:
            print(f"加载数据失败: {e}")
            return None

# 测试代码
if __name__ == "__main__":
    collector = DataCollector()
    # 测试获取股票基本信息
    stock_info = collector.get_stock_basic('000001')
    print("股票基本信息:", stock_info)
    
    # 测试获取日K线数据
    daily_data = collector.get_daily_data('000001')
    print("日K线数据前5行:\n", daily_data.head())